import cv2
import numpy as np
from rknn.api import RKNN

class sFaceFeatureExtractor:
    """人脸特征提取模型，支持传入原始图片"""
    
    def __init__(self, input_size=(112, 112), target='rk3588'):
        self.model_path = "./model/face_recognition_sface_2021dec.rknn"
        self.input_size = input_size
        self.target = target
        self.rknn = RKNN()
        self._load_model()
    
    def _load_model(self):
        if self.target != None:
            ret = self.rknn.load_rknn(self.model_path)
            if ret != 0:
                raise Exception(f"加载模型失败: {self.model_path}")    
            ret = self.rknn.init_runtime(target=self.target)
            if ret != 0:
                raise Exception("初始化运行环境失败")
        else:
            # 配置模型构建参数
            ret = self.rknn.config(mean_values=[0.0, 0.0, 0.0], std_values=[1.0, 1.0, 1.0], quantized_dtype='w8a8', target_platform='rk3588')
            if ret != 0:
                print('配置模型失败')
                exit(ret)
            # 加载原始模型（如ONNX）
            ret = self.rknn.load_onnx(model='./model/face_recognition_sface_2021dec.onnx')
            if ret != 0:
                print('加载ONNX模型失败')
                exit(ret)
            # 构建模型（指定target_platform为模拟器）
            ret = self.rknn.build(do_quantization=True, dataset='./model/dataset_RetinaFace.txt')
            if ret != 0:
                print('构建模型失败')
                exit(ret)
            # 初始化运行时环境（不指定target即为模拟器）
            ret = self.rknn.init_runtime()
            if ret != 0:
                print('初始化运行时环境失败')
                exit(ret)
            # Export rknn model
            print('--> Export rknn model')
            ret = self.rknn.export_rknn("./model/face_recognition_sface_2021dec.rknn")
            if ret != 0:
                print('Export rknn model failed!')
                exit(ret)
            print('done')
    
    def extract(self, image, face_boxes):
        """
        从多个人脸区域批量提取特征
        
        Args:
            image: numpy数组，BGR格式图像
            face_boxes: 人脸框坐标列表，每个框为 [x1, y1, x2, y2]
        Returns:
            features: 人脸特征向量列表，对应每个输入框
        """
        #if image is None or image.size == 0 or not face_boxes:
        #    return []
            
        valid_boxes = []
        processed_faces = []
        # 预处理所有面部框
        for box in face_boxes:
            x1, y1, x2, y2 = map(int, box[:4])
            
            # 确保人脸框在图像范围内
            x1 = max(0, x1)
            y1 = max(0, y1)
            x2 = min(image.shape[1], x2)
            y2 = min(image.shape[0], y2)
            
            if x1 >= x2 or y1 >= y2:
                processed_faces.append(None)  # 无效框添加None占位
                continue
                
            # 裁剪人脸区域
            face_roi = image[y1:y2, x1:x2]
            
            # 预处理
            processed_face = self._preprocess(face_roi)
            processed_faces.append(processed_face)
            valid_boxes.append(box)
        
        # 批量推理（逐个处理，RKNN可能不支持真正的批量）
        features = []
        for i, processed_face in enumerate(processed_faces):
            if processed_face is None:
                features.append(None)  # 无效框对应None特征
                continue
                
            # 模型推理
            try:
                outputs = self.rknn.inference(inputs=[processed_face])
                feature = outputs[0].flatten()
                features.append(feature)
            except Exception as e:
                print(f"特征提取失败，框索引 {i}: {str(e)}")
                features.append(None)
                
        return features
    
    def _preprocess(self, face_roi):
        """预处理人脸图像"""
        # 缩放至模型输入尺寸
        face_roi = cv2.resize(face_roi, self.input_size)
        
        # BGR转RGB
        face_roi = cv2.cvtColor(face_roi, cv2.COLOR_BGR2RGB)
        
        # 归一化
        #face_roi = face_roi.astype(np.float32) / 255.0
        
        # 调整维度为模型输入格式 [1, H, W, C]
        face_roi = np.expand_dims(face_roi, axis=0)
        
        return face_roi
    
    def release(self):
        if self.rknn:
            self.rknn.release()
